#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
小米便签网页版导出工具（自动ChromeDriver管理版）
该版本使用webdriver-manager自动管理ChromeDriver，无需手动下载和配置
"""

import os
import time
import json
import csv
import logging
import argparse
import re
from datetime import datetime
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException, NoSuchElementException
from webdriver_manager.chrome import ChromeDriverManager
import openpyxl
from openpyxl.styles import Alignment

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('mi_note_web_auto_driver.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

class XiaomiWebNoteExporter:
    def __init__(self, chrome_path=None, headless=False, timeout=30, export_format='excel', output_dir='./exports'):
        """初始化导出器"""
        self.chrome_path = chrome_path
        self.headless = headless
        self.timeout = timeout
        self.export_format = export_format.lower()
        self.output_dir = output_dir
        self.driver = None
        
        # 确保输出目录存在
        os.makedirs(self.output_dir, exist_ok=True)
    
    def _init_driver(self):
        """初始化WebDriver，使用webdriver-manager自动管理ChromeDriver"""
        logger.info("正在初始化Chrome浏览器...")
        
        chrome_options = Options()
        
        # 设置Chrome浏览器路径（如果提供）
        if self.chrome_path:
            logger.info(f"使用指定的Chrome浏览器路径: {self.chrome_path}")
            chrome_options.binary_location = self.chrome_path
        
        # 配置无头模式
        if self.headless:
            logger.info("启用无头模式")
            chrome_options.add_argument('--headless')
            chrome_options.add_argument('--disable-gpu')
            chrome_options.add_argument('--window-size=1920,1080')
        
        # 添加其他优化参数
        chrome_options.add_argument('--no-sandbox')
        chrome_options.add_argument('--disable-dev-shm-usage')
        chrome_options.add_argument('--start-maximized')
        chrome_options.add_argument('--disable-infobars')
        chrome_options.add_argument('--disable-extensions')
        chrome_options.add_argument('--ignore-certificate-errors')
        
        try:
            # 使用webdriver-manager自动下载和管理ChromeDriver
            logger.info("正在自动下载和配置匹配的ChromeDriver...")
            service = Service(ChromeDriverManager().install())
            self.driver = webdriver.Chrome(service=service, options=chrome_options)
            self.driver.set_page_load_timeout(self.timeout)
            self.driver.implicitly_wait(5)
            logger.info("Chrome浏览器初始化成功")
        except Exception as e:
            logger.error(f"Chrome浏览器初始化失败: {str(e)}")
            raise
    
    def login(self):
        """登录小米便签网页版"""
        logger.info("打开小米便签网页版...")
        self.driver.get('https://i.mi.com/note/')
        
        logger.info("请在浏览器中完成登录操作")
        logger.info("登录成功后请等待自动继续...")
        
        try:
            # 等待用户登录完成，通过检测便签列表区域是否出现
            WebDriverWait(self.driver, 120).until(
                EC.presence_of_element_located((By.CLASS_NAME, 'note-list'))
            )
            logger.info("登录成功，已进入便签列表页面")
        except TimeoutException:
            logger.error("登录超时，请检查网络连接或手动操作")
            raise
    
    def _get_note_content_from_right_panel(self):
        """从右侧面板获取便签内容"""
        try:
            # 检查右侧面板是否存在
            right_panel = WebDriverWait(self.driver, 5).until(
                EC.presence_of_element_located((By.CLASS_NAME, 'editor-wrapper'))
            )
            
            # 获取标题
            try:
                title_element = right_panel.find_element(By.CLASS_NAME, 'note-title')
                title = title_element.text.strip()
            except NoSuchElementException:
                title = "无标题"
                logger.warning("未找到便签标题，使用默认标题")
            
            # 获取内容
            try:
                content_element = right_panel.find_element(By.CLASS_NAME, 'note-content')
                # 获取所有段落文本
                paragraphs = content_element.find_elements(By.TAG_NAME, 'p')
                content = '\n'.join([p.text.strip() for p in paragraphs if p.text.strip()])
                
                # 如果没有段落，尝试获取整个元素的文本
                if not content:
                    content = content_element.text.strip()
            except NoSuchElementException:
                content = ""
                logger.warning("未找到便签内容")
            
            # 获取创建时间（如果可用）
            create_time = "未知"
            
            return {
                'title': title,
                'content': content,
                'create_time': create_time
            }
        except Exception as e:
            logger.error(f"从右侧面板获取便签内容失败: {str(e)}")
            return {
                'title': "获取失败",
                'content': f"错误: {str(e)}",
                'create_time': "未知"
            }
    
    def _get_note_content_from_list(self, note_element):
        """从列表项直接获取便签内容（备用方法）"""
        try:
            # 尝试获取标题
            try:
                title_element = note_element.find_element(By.CLASS_NAME, 'note-item-title')
                title = title_element.text.strip()
            except NoSuchElementException:
                title = "无标题"
            
            # 尝试获取内容预览
            try:
                content_element = note_element.find_element(By.CLASS_NAME, 'note-item-content')
                content = content_element.text.strip()
            except NoSuchElementException:
                content = ""
            
            return {
                'title': title,
                'content': content,
                'create_time': "未知"
            }
        except Exception as e:
            logger.error(f"从列表项获取便签内容失败: {str(e)}")
            return {
                'title': "获取失败",
                'content': f"错误: {str(e)}",
                'create_time': "未知"
            }
    
    def get_all_notes(self):
        """获取所有便签"""
        logger.info("开始获取所有便签...")
        notes = []
        
        try:
            # 等待便签列表加载完成
            WebDriverWait(self.driver, self.timeout).until(
                EC.presence_of_all_elements_located((By.CLASS_NAME, 'note-item'))
            )
            
            # 获取所有便签项
            note_elements = self.driver.find_elements(By.CLASS_NAME, 'note-item')
            logger.info(f"找到 {len(note_elements)} 条便签")
            
            # 遍历所有便签
            for index, note_element in enumerate(note_elements):
                try:
                    # 点击便签项，使其在右侧面板显示
                    note_element.click()
                    # 等待内容加载
                    time.sleep(0.5)
                    
                    # 从右侧面板获取内容
                    note_content = self._get_note_content_from_right_panel()
                    
                    # 如果右侧面板内容为空，尝试从列表项直接获取
                    if not note_content['content'] or note_content['content'].startswith('错误:'):
                        note_content = self._get_note_content_from_list(note_element)
                    
                    # 添加便签到列表
                    notes.append(note_content)
                    logger.info(f"获取便签 {index + 1}/{len(note_elements)}: {note_content['title']}")
                    
                except Exception as e:
                    logger.error(f"获取第 {index + 1} 条便签失败: {str(e)}")
                    # 添加一个错误便签，继续处理其他便签
                    notes.append({
                        'title': f"便签 {index + 1} 获取失败",
                        'content': f"错误: {str(e)}",
                        'create_time': "未知"
                    })
            
            logger.info(f"成功获取 {len(notes)} 条便签")
            return notes
        except Exception as e:
            logger.error(f"获取便签列表失败: {str(e)}")
            raise
    
    def export_notes(self, notes):
        """导出便签到指定格式"""
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        filename = os.path.join(self.output_dir, f'mi_notes_{timestamp}')
        
        try:
            if self.export_format == 'excel':
                self._export_to_excel(notes, filename + '.xlsx')
            elif self.export_format == 'csv':
                self._export_to_csv(notes, filename + '.csv')
            elif self.export_format == 'json':
                self._export_to_json(notes, filename + '.json')
            elif self.export_format == 'txt':
                self._export_to_text(notes, filename + '.txt')
            else:
                logger.error(f"不支持的导出格式: {self.export_format}")
                raise ValueError(f"不支持的导出格式: {self.export_format}")
            
            logger.info(f"便签已成功导出到: {filename}.{self.export_format}")
            return filename + f".{self.export_format}"
        except Exception as e:
            logger.error(f"导出便签失败: {str(e)}")
            raise
    
    def _export_to_excel(self, notes, filename):
        """导出便签到Excel文件"""
        wb = openpyxl.Workbook()
        ws = wb.active
        ws.title = "小米便签"
        
        # 设置表头
        headers = ['标题', '内容', '创建时间']
        for col_idx, header in enumerate(headers, 1):
            ws.cell(row=1, column=col_idx, value=header)
            # 设置表头样式
            ws.cell(row=1, column=col_idx).alignment = Alignment(horizontal='center', vertical='center')
        
        # 填充数据
        for row_idx, note in enumerate(notes, 2):
            ws.cell(row=row_idx, column=1, value=note['title'])
            ws.cell(row=row_idx, column=2, value=note['content'])
            ws.cell(row=row_idx, column=3, value=note['create_time'])
            
            # 设置内容自动换行
            ws.cell(row=row_idx, column=2).alignment = Alignment(wrap_text=True)
        
        # 调整列宽
        ws.column_dimensions['A'].width = 30
        ws.column_dimensions['B'].width = 80
        ws.column_dimensions['C'].width = 20
        
        wb.save(filename)
    
    def _export_to_csv(self, notes, filename):
        """导出便签到CSV文件"""
        with open(filename, 'w', newline='', encoding='utf-8-sig') as f:
            writer = csv.writer(f)
            # 写入表头
            writer.writerow(['标题', '内容', '创建时间'])
            # 写入数据
            for note in notes:
                writer.writerow([note['title'], note['content'], note['create_time']])
    
    def _export_to_json(self, notes, filename):
        """导出便签到JSON文件"""
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(notes, f, ensure_ascii=False, indent=2)
    
    def _export_to_text(self, notes, filename):
        """导出便签到文本文件"""
        with open(filename, 'w', encoding='utf-8') as f:
            for i, note in enumerate(notes, 1):
                f.write(f"===== 便签 {i}: {note['title']} =====\n")
                f.write(f"{note['content']}\n")
                if note['create_time'] != '未知':
                    f.write(f"\n创建时间: {note['create_time']}\n")
                f.write("\n" + "="*50 + "\n\n")
    
    def run(self):
        """运行导出流程"""
        try:
            # 初始化浏览器
            self._init_driver()
            
            # 登录
            self.login()
            
            # 获取所有便签
            notes = self.get_all_notes()
            
            # 导出便签
            if notes:
                export_file = self.export_notes(notes)
                logger.info(f"导出完成！共导出 {len(notes)} 条便签到 {export_file}")
            else:
                logger.warning("没有找到任何便签")
            
        except Exception as e:
            logger.error(f"导出过程中发生错误: {str(e)}")
            raise
        finally:
            # 关闭浏览器
            if self.driver:
                logger.info("正在关闭浏览器...")
                self.driver.quit()
                logger.info("浏览器已关闭")
    

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='小米便签网页版导出工具（自动ChromeDriver管理版）')
    parser.add_argument('--chrome_path', type=str, help='Chrome浏览器可执行文件路径（可选）')
    parser.add_argument('--headless', action='store_true', help='启用无头模式运行')
    parser.add_argument('--timeout', type=int, default=30, help='超时时间（秒），默认为30秒')
    parser.add_argument('--format', type=str, choices=['excel', 'csv', 'json', 'txt'], default='excel', 
                        help='导出格式，支持excel、csv、json、txt，默认为excel')
    parser.add_argument('--output_dir', type=str, default='./exports', help='导出文件保存目录，默认为./exports')
    
    return parser.parse_args()


def main():
    """主函数"""
    try:
        # 解析命令行参数
        args = parse_arguments()
        
        logger.info(f"小米便签网页版导出工具启动")
        logger.info(f"导出格式: {args.format}")
        logger.info(f"输出目录: {args.output_dir}")
        logger.info(f"无头模式: {args.headless}")
        logger.info(f"超时时间: {args.timeout}秒")
        if args.chrome_path:
            logger.info(f"Chrome浏览器路径: {args.chrome_path}")
        
        # 创建导出器实例
        exporter = XiaomiWebNoteExporter(
            chrome_path=args.chrome_path,
            headless=args.headless,
            timeout=args.timeout,
            export_format=args.format,
            output_dir=args.output_dir
        )
        
        # 运行导出流程
        exporter.run()
        
        logger.info("小米便签网页版导出工具运行完成")
    except Exception as e:
        logger.error(f"程序运行失败: {str(e)}")
        print(f"错误: {str(e)}")
        exit(1)


if __name__ == '__main__':
    main()